Summary

Adds LanceStream, a streaming IVF kmeans training path that keeps peak raw-vector training memory bounded for very large IVF partition counts while preserving the existing non-streaming trainer when streaming is not requested.

The latest update also bounds the sampling read buffer: fixed sampled ranges are split into smaller range reads and take_scan readahead defaults to one batch. This avoids the previous RSS spike from large contiguous sampled ranges while keeping an environment override for benchmark tuning.

Feature

• Adds streaming_sample_rate, streaming_coreset_rate, and streaming_refine_passes to IVF build parameters.
• Exposes the new parameters through the Python index creation path.
• Keeps num_partitions > 256 on hierarchical kmeans; large-k streaming does not fall back to flat kmeans.
• Uses fixed sampled row ranges for streaming training so coreset construction and optional refine passes read the same sample.
• Splits fixed sample ranges into bounded sub-ranges before take_scan; default range chunk is 8192 rows.
• Adds LANCE_STREAMING_IVF_PREFETCH_DEPTH and LANCE_STREAMING_IVF_TAKE_RANGE_ROWS for experimental tuning. Default prefetch depth is 1 because deeper prefetch gave minimal speedup and higher RSS.

Algorithm

The existing IVF trainer samples up to num_partitions * sample_rate raw vectors before training. For very large num_partitions, that materializes too many raw vectors at once. LanceStream decouples total sample size from peak raw-vector memory.

The streaming path works as follows:

1. Build a fixed sample plan as sorted random row ranges, bounded by num_partitions * sample_rate and dataset row count.
2. Split those ranges into training chunks with at most num_partitions * streaming_sample_rate raw vectors.
3. For each chunk, load only the current raw vectors, train a local kmeans summary, assign raw vectors to local centroids, and record weighted local centroids plus local loss contribution.
4. Merge weighted local summaries into a global coreset bounded by num_partitions * streaming_coreset_rate.
5. Train final IVF centroids with weighted hierarchical kmeans over the coreset. The hierarchical trainer starts with a bounded fanout and repeatedly splits high-loss/high-weight partitions until it reaches num_partitions.
6. Run a short weighted Lloyd refinement over the coreset centroids.
7. Optionally run streaming_refine_passes raw-vector Lloyd passes over the same fixed sample. Each pass streams chunks, assigns raw vectors to current centroids, accumulates global sums/counts/loss, updates non-empty centroids at pass end, and preserves old centroids for empty clusters.

Peak raw-vector memory is bounded by:

num_partitions * streaming_sample_rate * dimension * sizeof(float32)

The coreset budget is separately bounded by:

num_partitions * streaming_coreset_rate

Centroid and accumulator buffers still scale with num_partitions * dimension.

Benchmarks

Setup

• VM: GCP c2-standard-16, 16 vCPU, ~64 GiB RAM, AVX512, 500 GB disk.
• Dataset: real SIFT vectors, 128-dimensional float32.
• Dataset choice: k=1024 on SIFT1M subset; k>=4096 on SIFT1B prefix.
• Training sample size: first k * 256 rows for unified exact-loss evaluation; training sample size also k * 256.
• LanceStream defaults in this run: stream=64, coreset=16, prefetch=1, refine=0, sample_rate=256, max_iters=50.
• Tuned LanceStream variants only changed stream, coreset, and prefetch_depth.
• Timeout: 1 hour per run.

Loss Methodology

All comparable losses below are computed outside the training algorithm:

1. Save trained centroids to f32 files.
2. Read the first k * 256 vectors from the dataset.
3. Compute exact kmeans objective with faiss.IndexFlatL2:

sum_x min_c ||x - c||^2

No table below uses algorithm-internal reported loss for comparison. Exact loss is reported for k <= 16,384. For k=65,536 and 131,072, exact loss was not run because one evaluation is roughly O(k * k * 256) distance work; those rows report train time/RSS/status only.

LanceStream Tuning

Percentages are relative to LanceStream stream=64, coreset=16, prefetch=1 at the same k.

Takeaways:

• Increasing prefetch_depth is not a good default. At 128K, p16 was only 1.6% faster than p1 but used 6.7% more RSS. At smaller k it often increased RSS much more.
• Lowering coreset from 16 to 8 is the strongest speed/RSS knob, but it can raise loss. At 16K it was 38.7% faster and 12.4% lower RSS, but loss was 6.0% worse.
• Smaller stream lowers RSS but does not necessarily improve time. At 128K, stream=16, coreset=8 used only 2.39 GiB but took 17.0 minutes.

Algorithm Comparison

Percentages are relative to LanceStream stream=64, coreset=16, prefetch=1 at the same k.

Large-k notes:

• Faiss and MiniBatch both timed out at 16K, so 64K and 128K were skipped.
• BICO-style completed 16K but took 45.1 minutes and timed out at 64K.
• treeCoreset completed 16K but used 12.32 GiB RSS; at 64K it crashed with signal 11 after reaching 22.31 GiB RSS.
• Lance non-stream is faster but scales memory with the full k * 256 raw sample; at 128K it used 32.62 GiB RSS on this 128D dataset.
• LanceStream completed every tested k with much lower RSS than Lance non-stream. At 128K, default LanceStream used 6.02 GiB; the lower-memory stream=16, coreset=8 variant used 2.39 GiB.

Validation

• cargo fmt --all --check
• cargo check -p lance --locked
• cargo test -p lance --lib test_split_ranges_by_row_count --locked on a clean temp worktree containing only the production ivf.rs change
• git diff --check